[Fix][Relax][ONNX] Import TopK indices as int64#19973
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the ONNX frontend's TopK operator implementations (_impl_v11 and _impl_v1) to explicitly specify dtype="int64" when calling relax.op.topk. Additionally, the corresponding test_topk test is updated to enable data type checking. I have no feedback to provide as there are no review comments.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
cc @guan404ming. 🙌 |
80cffed to
319a74b
Compare
ONNX specifies that TopK's second output, indices, has int64 element type. The Relax ONNX frontend previously used relax.op.topk without specifying the indices dtype, so Relax used its default int32 indices. This can make valid ONNX graphs fail or produce an incorrect imported type when TopK indices are consumed by later integer or indexing operations. Pass dtype="int64" when importing ONNX TopK and enable dtype checking in the TopK frontend test. Signed-off-by: viiccwen <vicwen@apache.org>
2db4420 to
0a140b4
Compare
|
LGTM but the ci seems not worked for now. cc @tqchen |
|
Thx @guan404ming, @tlopex for reviewing! |
4 participants
Fixes #19972
ONNX specifies that the second output of TopK,
indices, has element typeint64, and the ONNX TopK operator spec constrains the index tensor type totensor(int64): https://onnx.ai/onnx/operators/onnx__TopK.htmlThe Relax ONNX frontend previously called
relax.op.topkwithout specifying the output indices dtype, so Relax used its defaultint32indices.This can make otherwise valid ONNX graphs fail during import when the TopK indices are consumed by later integer/index operations that use ONNX's usual
int64constants. One example isTopK -> Div, where Relax rejects the binary operation because the imported TopK indices areint32while the divisor isint64.This patch passes
dtype="int64"when importing ONNX TopK, matching the ONNX operator spec. It also updates the existing TopK frontend test to check output dtypes, so the imported indices must match ONNX Runtime'sint64output.Verification:
uv run --no-sync python -m pytest tests/python/relax/test_frontend_onnx.py::test_topk -q